#!/bin/python

import numpy as np
import cPickle as pickle

def gaussian_pulse(height, y_offset, st_dev, times):
    """
    Raised windowed Gaussian for locally-optimal driving
    """
    mean = np.amax(times)/2.
    gaus_array = np.exp(-(times - mean)**2/(2. * st_dev ** 2)) + y_offset
    return height * gaus_array/np.amax(gaus_array)

input_dict = dict()

#input_dict['Navg'] = 10                           # number of trajectories
input_dict['dim'] = 8
input_dict['phi'] = 0.5*np.pi                       #phase of LO

input_dict['eta'] = 0.5
Delta_q = 1600     
g = np.sqrt(Delta_q)              # corresponds to chi = 1
input_dict['chi'] = g**2/Delta_q
input_dict['kappa1'] = 2*input_dict['chi']
input_dict['kappa2'] = 2*input_dict['chi']
input_dict['Delta_a'] = np.sqrt(3)*input_dict['chi']
input_dict['Delta_b'] = -np.sqrt(3)*input_dict['chi']
input_dict['gamma_d1'] = 0./300.
input_dict['gamma_d2'] = 0./300.
input_dict['gamma_d3'] = 0./300.
input_dict['gamma_1'] = 0.0
input_dict['tmeas_on'] = 0./input_dict['chi']
input_dict['tmeas_off'] = 5./(input_dict['chi']*input_dict['eta'])

ton = input_dict['tmeas_on'] 
toff = input_dict['tmeas_off'] - 2

distortion_fudge_factor = 1.2 #In order to account for all the distortion, 
                              #we have to integrate for additional time.

input_dict['T'] = distortion_fudge_factor * input_dict['tmeas_off'] # final time
input_dict['dt'] = 0.005                     # time step
input_dict['N'] = int(round(input_dict['T']/input_dict['dt']))                       # number of time steps


#epsilon0 = 2.0/np.sqrt(toff - ton) # This is to get the error probability = 0.025
epsilon0 = 2.0/np.sqrt(16) # This is to get the error probability = 0.025

sigma = 10

# Initial state

initstate = 1./np.sqrt(input_dict['dim'])*np.ones((input_dict['dim'],1), dtype='complex128')
initMat = initstate*initstate.T

input_dict['init'] = np.reshape(initMat.T,input_dict['dim']**2,1)

#ideal_p = 0.5*[1,0,0,1,0,1,1,0]
#ideal_p = sparse(ideal_p.T*ideal_p)
#ideal_m = 0.5*[0,1,1,0,1,0,0,1]
#ideal_m = sparse(ideal_m.T*ideal_m)

#rho_save = zeros(dim,dim,Navg)

#Time-independent part of Liouvillian
sigmaZ = np.array([[1, 0], [0, - 1]])
sigmaM = np.array([[0, 1], [0, 0]])
eye = np.eye(2)
input_dict['fulleye'] = np.eye(input_dict['dim'])

# The dephasing operator for qubit 1
Ad1 = reduce(np.kron,(sigmaZ, eye, eye))
Ld1 = np.kron(Ad1, Ad1.T) - np.dot(0.5, np.kron(np.dot(
    Ad1.T, Ad1), input_dict['fulleye'])) - np.dot(0.5, np.kron(input_dict['fulleye'], np.dot(Ad1.T, Ad1).T))

# The dephasing operator for qubit 2
Ad2 = reduce(np.kron,(eye, sigmaZ, eye))
Ld2 = np.kron(Ad2, Ad2.T) - np.dot(0.5, np.kron(np.dot(
    Ad2.T, Ad2), input_dict['fulleye'])) - np.dot(0.5, np.kron(input_dict['fulleye'], np.dot(Ad2.T, Ad2).T))

# The dephasing operator for qubit 3
Ad3 = reduce(np.kron,(eye, eye, sigmaZ))
Ld3 = np.kron(Ad3, Ad3.T) - np.dot(0.5, np.kron(np.dot(
    Ad3.T, Ad3), input_dict['fulleye'])) - np.dot(0.5, np.kron(input_dict['fulleye'], np.dot(Ad3.T, Ad3).T))

# The relaxation operator for qubit 1
Ar1 = reduce(np.kron,(sigmaM, eye, eye))
Ar1dagger = Ar1.T
Lr1 = np.kron(Ar1, Ar1dagger.T) - np.dot(0.5, np.kron(np.dot(
    Ar1dagger, Ar1), input_dict['fulleye'])) - np.dot(0.5, np.kron(input_dict['fulleye'], np.dot(Ar1dagger, Ar1).T))

# The relaxation operator for qubit 2
Ar2 = reduce(np.kron,(eye, sigmaM, eye))
Ar2dagger = Ar2.T
Lr2 = np.kron(Ar2, Ar2dagger.T) - np.dot(0.5, np.kron(np.dot(
    Ar2dagger, Ar2), input_dict['fulleye'])) - np.dot(0.5, np.kron(input_dict['fulleye'], np.dot(Ar2dagger, Ar2).T))

# The relaxation operator for qubit 3
Ar3 = reduce(np.kron,(eye, eye, sigmaM))
Ar3dagger = Ar3.T
Lr3 = np.kron(Ar3, Ar3dagger.T) - np.dot(0.5, np.kron(np.dot(
    Ar3dagger, Ar3), input_dict['fulleye'])) - np.dot(0.5, np.kron(input_dict['fulleye'], np.dot(Ar3dagger, Ar3).T))

# The time independent deterministic part of Liouvillian
input_dict['D1'] = (np.dot(np.dot(0.5, input_dict['gamma_d1']), Ld1) + \
					np.dot(np.dot(0.5, input_dict['gamma_d2']), Ld2) + \
					np.dot(np.dot(0.5, input_dict['gamma_d3']), Ld3) + \
					np.dot(np.dot(0, (np.dot((g / Delta_q) ** 2, (input_dict['kappa1'] + input_dict['kappa2'])) + \
						input_dict['gamma_1'])), (Lr1 + Lr2 + Lr3))).astype(np.complex128)

# Calculate the fields in the cavity
# TODO if slow: replace with analytic result
input_dict['tvec'] = np.linspace(0, input_dict['T'], input_dict['N'])

input_dict['epsilon'] = np.dot(epsilon0 / np.pi, (np.arctan(np.dot(sigma, (input_dict['tvec'] - ton))) + np.dot(
    0.5, np.pi))) - np.dot(epsilon0 / np.pi, (np.arctan(np.dot(sigma, (input_dict['tvec'] - toff))) + np.dot(0.5, np.pi)))

#Ideal Plus/Minus vectorized density matrices
ideal_p = 0.5 * np.array([1,0,0,1,0,1,1,0], dtype=np.complex128).T
input_dict['ideal_p'] = np.kron(ideal_p, ideal_p)
ideal_m = 0.5 * np.array([0,1,1,0,1,0,0,1], dtype=np.complex128).T
input_dict['ideal_m'] = np.kron(ideal_m, ideal_m)

with open('trajectory_input_file', 'w') as file:
    pickle.dump(input_dict, file)
